--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
--
--  file:    unit_smart_select.lua
--  version: 1.3
--  brief:   Selects units as you drag over them and provides selection modifier hotkeys
--  original author: Ryan Hileman (aegis)
--
--  Copyright (C) 2011.
--  Public Domain.
--
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

function widget:GetInfo()
	return {
		name      = "SmartSelect",
		desc      = "Selects units as you drag over them and provides selection modifier hotkeys",
		author    = "aegis",
		date      = "Jan 2, 2011",
		license   = "Public Domain",
		layer     = 0,
		enabled   = true
	}
end

-----------------------------------------------------------------
-- user config section
-----------------------------------------------------------------

-- whether to select buildings when mobile units are inside selection rectangle
local selectBuildingsWithMobile = false

-- only select new units identical to those already selected
local sameSelectKey = 'z'

-----------------------------------------------------------------
-- manually generated locals because I don't have trepan's script
-----------------------------------------------------------------
local GetTimer = Spring.GetTimer
local DiffTimers = Spring.DiffTimers
local GetMouseState = Spring.GetMouseState
local GetModKeyState = Spring.GetModKeyState
local GetKeyState = Spring.GetKeyState

local TraceScreenRay = Spring.TraceScreenRay
local WorldToScreenCoords = Spring.WorldToScreenCoords

local GetMyTeamID = Spring.GetMyTeamID
local GetMyPlayerID = Spring.GetMyPlayerID
local GetPlayerInfo = Spring.GetPlayerInfo
local GetTeamUnits = Spring.GetTeamUnits
local GetAllUnits = Spring.GetAllUnits
local GetSelectedUnits = Spring.GetSelectedUnits
local GetUnitsInRectangle = Spring.GetUnitsInRectangle
local SelectUnitArray = Spring.SelectUnitArray
local GetActiveCommand = Spring.GetActiveCommand

local GetGroundHeight = Spring.GetGroundHeight
local GetMiniMapGeometry = Spring.GetMiniMapGeometry
local GetMiniMapDualScreen = Spring.GetMiniMapDualScreen
local IsAboveMiniMap = Spring.IsAboveMiniMap

local GetUnitDefID = Spring.GetUnitDefID
local GetUnitPosition = Spring.GetUnitPosition

local UnitDefs = UnitDefs
local min = math.min
local max = math.max

local glColor = gl.Color
local glVertex = gl.Vertex
local glLineWidth = gl.LineWidth
local glDepthTest = gl.DepthTest
local glBeginEnd = gl.BeginEnd
local GL_LINE_STRIP = GL.LINE_STRIP
-----------------------------------------------------------------
-- end function locals ------------------------------------------
-----------------------------------------------------------------

sameSelectKey = Spring.GetKeyCode(sameSelectKey)
local minimapOnLeft = (Spring.GetMiniMapDualScreen() == "left")

local combatFilter, builderFilter, buildingFilter, mobileFilter

local referenceCoords
local referenceScreenCoords
local referenceSelection
local referenceSelectionTypes
local lastSelection
local minimapRect

local lastCoords
local lastMeta

local lastUpdate = GetTimer()
local spectating = false
local myPlayerID

local function sort(v1, v2)
	if v1 > v2 then
		return v2, v1
	else
		return v1, v2
	end
end

local function GetUnitsInArbitraryRectangle(x1, z1, x2, z2, team)
	x1, x2 = sort(x1, x2)
	z1, z2 = sort(z1, z2)

	local units = GetUnitsInRectangle(x1, z1, x2, z2, team)
	return units
end


local mapWidth, mapHeight = Game.mapSizeX, Game.mapSizeZ
local function MinimapToWorldCoords(x, y)
	px, py, sx, sy = GetMiniMapGeometry()
	local plx = 0
	if (minimapOnLeft) then plx = sx end

	x = ((x - px + plx) / sx) * mapWidth
	local z = (1 - ((y - py) / sy)) * mapHeight
	y = GetGroundHeight(x, z)
	return x, y, z
end

local function GetUnitsInMinimapRectangle(x1, y1, x2, y2, team)
	left, _, top = MinimapToWorldCoords(x1, y1)
	right, _, bottom = MinimapToWorldCoords(x2, y2)

	local left, right = sort(left, right)
	local bottom, top = sort(bottom, top)

	minimapRect = {left, bottom, right, top}
	return GetUnitsInRectangle(left, bottom, right, top, team)
end

local function GetUnitsInScreenRectangle(x1, y1, x2, y2, team)
	local units
	if (team) then
		units = GetTeamUnits(team)
	else
		units = GetAllUnits()
	end
	
	local left, right = sort(x1, x2)
	local bottom, top = sort(y1, y2)

	local result = {}

	for i=1, #units do
		local uid = units[i]
		x, y, z = GetUnitPosition(uid)
		x, y = WorldToScreenCoords(x, y, z)
		if (left <= x and x <= right) and (top >= y and y >= bottom) then
			result[#result+1] = uid
		end
	end
	return result
end

function widget:MousePress(x, y, button)
	if (button == 1) then
		referenceSelection = GetSelectedUnits()
		referenceSelectionTypes = {}
		for i=1, #referenceSelection do
			udid = GetUnitDefID(referenceSelection[i])
			referenceSelectionTypes[udid] = 1
		end
		referenceScreenCoords = {x, y}
		lastMeta = nil
		lastSelection = nil

		if (IsAboveMiniMap(x, y)) then
			referenceCoords = {0, 0, 0}
			lastCoords = {0, 0, 0}
		else
			local _, c = TraceScreenRay(x, y, true, false, true)

			referenceCoords = c
			lastCoords = c
		end
	end
end

function widget:Update()
	--[[
	local newUpdate = GetTimer()
	if (DiffTimers(newUpdate, lastUpdate) < 0.1) then
		return
	end
	lastUpdate = newUpdate
	--]]

	if (referenceCoords ~= nil and GetActiveCommand() == 0) then
		x, y, pressed = GetMouseState()
		local px, py, sx, sy = GetMiniMapGeometry()
		
		if (pressed) and (referenceSelection ~= nil) then
			local alt, ctrl, meta, shift = GetModKeyState()
			if (#referenceSelection == 0) then
				-- no point in inverting an empty selection
				ctrl = false
			end

			local sameSelect = GetKeyState(sameSelectKey)
			
			local sameLast = (referenceScreenCoords ~= nil) and (x == referenceScreenCoords[1] and y == referenceScreenCoords[2])
			if (sameLast and lastCoords == referenceCoords) then
				return
			end

			if sameLast and (lastMeta ~= nil) and (alt == lastMeta[1] and ctrl == lastMeta[2]
						and meta == lastMeta[3] and shift == lastMeta[4])
				then return end

			lastCoords = {x, y}
			lastMeta = {alt, ctrl, meta, shift}

			local mouseSelection, originalMouseSelection
			local r = referenceScreenCoords
			local playing = GetPlayerInfo(myPlayerID).spectating == false
			local team = (playing and GetMyTeamID())
			if (r ~= nil and IsAboveMiniMap(r[1], r[2])) then
				local mx, my = max(px, min(px+sx, x)), max(py, min(py+sy, y))
				mouseSelection = GetUnitsInMinimapRectangle(r[1], r[2], x, y, team)
			else
				local d = referenceCoords
				local x1, y1 = WorldToScreenCoords(d[1], d[2], d[3])
				mouseSelection = GetUnitsInScreenRectangle(x, y, x1, y1, team)
			end
			
			local newSelection = {}
			local uid, udid, udef, tmp

			if (sameSelect) and (#referenceSelection > 0) then
				-- only select new units identical to those already selected
				tmp = {}
				for i=1, #mouseSelection do
					uid = mouseSelection[i]
					udid = GetUnitDefID(uid)
					if (referenceSelectionTypes[udid] ~= nil) then
						tmp[#tmp+1] = uid
					end
				end
				mouseSelection = tmp
			end

			if (alt) then
				-- only select mobile combat units

				if (ctrl == false) then
					tmp = {}
					for i=1, #referenceSelection do
						uid = referenceSelection[i]
						udid = GetUnitDefID(uid)
						if (combatFilter[udid]) then -- is a combat unit
							tmp[#tmp+1] = uid
						end
					end
					newSelection = tmp
				end

				tmp = {}
				for i=1, #mouseSelection do
					uid = mouseSelection[i]
					udid = GetUnitDefID(uid)
					if (combatFilter[udid]) then -- is a combat unit
						tmp[#tmp+1] = uid
					end
				end
				mouseSelection = tmp
			elseif (selectBuildingsWithMobile == false) and (shift == false) and (ctrl == false) then
				-- only select mobile units, not buildings
				local mobiles = false
				for i=1, #mouseSelection do
					uid = mouseSelection[i]
					udid = GetUnitDefID(uid)
					if (mobileFilter[udid]) then
						mobiles = true
						break
					end
				end

				if (mobiles) then
					tmp = {}
					for i=1, #mouseSelection do
						uid = mouseSelection[i]
						udid = GetUnitDefID(uid)
						if (buildingFilter[udid] == false) then
							tmp[#tmp+1] = uid
						end
					end
					mouseSelection = tmp
				end
			end

			if (#newSelection < 1) then
				newSelection = referenceSelection
			end

			if (ctrl) then
				-- deselect units inside the selection rectangle, if we already had units selected
				local negative = {}

				for i=1, #mouseSelection do
					uid = mouseSelection[i]
					negative[uid] = 1
				end

				tmp = {}
				for i=1, #newSelection do
					uid = newSelection[i]
					if (negative[uid] == nil) then
						tmp[#tmp+1] = uid
					end
				end
				newSelection = tmp
				SelectUnitArray(newSelection)
			elseif (shift) then
				--Spring.Echo('shift')
				-- append units inside selection rectangle to current selection
				SelectUnitArray(newSelection)
				SelectUnitArray(mouseSelection, true)
			elseif (#mouseSelection > 0) then
				-- select units inside selection rectangle
				SelectUnitArray(mouseSelection)
			elseif (alt == false) then
				-- keep current selection while dragging until more things are selected
				SelectUnitArray(referenceSelection)
				lastSelection = nil
				return
			end

			lastSelection = GetSelectedUnits()
		elseif (lastSelection ~= nil) then
			SelectUnitArray(lastSelection)
			lastSelection = nil
			referenceSelection = nil
			referenceSelectionTypes = nil
			referenceCoords = nil
			minimapRect = nil
		else
			referenceSelection = nil
			referenceSelectionTypes = nil
			referenceCoords = nil
			minimapRect = nil
		end
	end
end

function widget:Initialize()
	myPlayerID = GetMyPlayerID()
	combatFilter = {}
	builderFilter = {}
	buildingFilter = {}
	mobileFilter = {}

	for udid, udef in pairs(UnitDefs) do
		local mobile = (udef.canMove and udef.speed > 0.000001)
		local builder = (udef.canReclaim and udef.reclaimSpeed > 0) or
						(udef.builder and udef.buildSpeed > 0) or
						(udef.canResurrect and udef.resurrectSpeed > 0) or
						(udef.canRepair and udef.repairSpeed > 0)
		local building = (mobile == false)
		local combat = (builder == false) and (mobile == true) and (#udef.weapons > 0)

		combatFilter[udid] = combat
		builderFilter[udid] = builder
		buildingFilter[udid] = building
		mobileFilter[udid] = mobile
	end
end

local function DrawRectangle(r)
	local x1, y1, x2, y2 = r[1], r[2], r[3], r[4]
	glVertex(r[1], 0, r[2])
	glVertex(r[1], 0, r[4])
	glVertex(r[3], 0, r[4])
	glVertex(r[3], 0, r[2])
	glVertex(r[1], 0, r[2])
end

function widget:DrawWorld()
	if (minimapRect ~= nil) then
		glColor(1, 1, 1, 1)
		glLineWidth(1.0)
		glDepthTest(false)

		glBeginEnd(GL_LINE_STRIP, DrawRectangle, minimapRect)
	end
end